import datetime
from LSTM.variableset import EMVar
import tushare as ts
from emutils import stk_utils
from sysdata import sys_base_data as sysBaseData
import pandas as pd
from emutils import time_utils
import traceback
from emmongo import em_mongomanager as mongoManager
from dateutil.parser import parse
from emmodels.em_position2 import Assets

class Rollback(object):

    MODE_MIN = 'MIN'
    MODE_DAY = 'DAY'


    def __init__(self, mode, is_real,assets_name, init_assets=50000,assets_detail='', strategy_name='', start_date='2010-01-01', end_date=EMVar.STR_NOW_YMD, params=None,**kwargs):
        self.start_date = start_date
        self.end_date = end_date
        self.mode = mode
        self.is_real = is_real

        sh_his_datas = ts.get_k_data(code=EMVar.CODE_SH, index=True, start=self.start_date, end=self.end_date)
        self.date_series = sh_his_datas[EMVar.date]
        if mode == Rollback.MODE_MIN:
            self.time_list = stk_utils.em_get_day_min_list()
        self.base_info_dic = {}
        self.total_tmp_min_data = pd.DataFrame()
        self.total_tmp_day_data = pd.DataFrame()
        self.registe_stks = []
        self.strategy_name = strategy_name
        self.append_now = False
        self.params = params

        self.assets = Assets(init_assets=init_assets, strategy_name=assets_name, strategy_detail=assets_detail, own_stock={})
        self.assets.start_date = self.start_date

        self.custom_init(**kwargs)
        pass

    def init_assets(self,):

        pass

    def draw_strategy_stks(self):
        # if date in self.total_datas.index:
        #     stks = self.total_datas[date:date][EMVar.code]
        #     for stk in stks:
        #         date_l, date_index_offset = stk_utils.em_get_date_offset(date, self.date_series, pre_offset=-5,
        #                                                                  later_offset=15)
        #         if date_l:
        #             self.draw_begin_date = date_l[0]
        #             self.draw_end_date = date_l[-1]
        #
        #             sfund = sysBaseData.get_self_fund(stk)
        #             df = sfund.get_index(start=self.start_date, append_now=False, autype='qfq')
        #             fileName = '%s_%s' % (date, sfund.code_with_name())
        #             draw_utils.draw_day_k(df[self.draw_begin_date: self.draw_end_date], isShow=False,
        #                               save_path=EMPath.strateg_info_file_full_path(is_common=True, strategyName=self.assets.strategy_name, fileName=fileName),
        #                               buy_date=date)
        #
        # return

        pass

    def custom_init(self, **kwargs):
        if 'recover_from_db' in kwargs.keys() and kwargs['recover_from_db']:
            self.assets.recover_from_db()
        pass

    def begin(self):
        if self.is_real:
            self.begin_real()
        else:
            if self.mode == self.__class__.MODE_DAY:
                for date in list(self.date_series):
                    try:
                        # 对缓存的分钟数据置空
                        self.__begin_day(date)
                    except Exception as e:
                        print(e)
                        traceback.print_tb(e.__traceback__)
                        continue
            elif self.mode == self.__class__.MODE_MIN:
                for date in list(self.date_series):
                    try:
                        date_time_list = [date + ' ' + x for x in self.time_list]
                        # 对缓存的分钟数据置空
                        self.base_info_dic = {}
                        self.before_open(date)
                        self.detect_stks(stks=self.registe_stks, date=date)
                        for date_time in date_time_list:
                            self.active(date_time)
                        self.after_close(date)
                    except Exception as e:
                        print(e)
                        continue



    def __begin_day(self, date):
        self.base_info_dic = {}
        self.before_open(date)
        self.detect_stks(stks=self.registe_stks, date=date)
        self.active(date)
        self.after_close(date)

    def begin_real(self, date=EMVar.format_now()):
        # 获取时间范围在监控范围内的stk
        self.mode = Rollback.MODE_DAY
        self.append_now = True

        # 对缓存的分钟数据置空
        self.__begin_day(date)


    def util_load_strategy_stks_start_end(self, date):
        stks = mongoManager.db[mongoManager.COL_SELF_FUND_SET].find(
            {'startDay': {'$lte': date}, 'endDay': {'$gte': date}, 'strategy_name': self.strategy_name}, {'stkNo': 1})
        if stks.count() > 0:
            l = list(stks)
            stks = [x['stkNo'] for x in l if len(x)]
            return stks
        return []

    def util_load_strategy_stks(self, begin_date, end_date):

        stks = mongoManager.db[mongoManager.COL_SELF_FUND_SET].find({'$and':[{'createDate': {'$lte': parse(end_date)}},
                                                                             {'createDate': {'$gte': parse(begin_date)}},
                                                                             {'strategy_name': self.strategy_name}]}, {'stkNo': 1})
        if stks.count() > 0:
            l = list(stks)
            stks = [x['stkNo'] for x in l if len(x)]
            return stks
        return []

    def __load_strategy_stks(self, date):
        return []

    def before_open(self, date):
        pass

    def after_close(self, date):
        pass

    def deliever(self, date_time, bar_datas):
        pass

    def active(self, date_time):
        if self.mode == self.__class__.MODE_MIN:
            if not self.total_tmp_min_data.empty:
                stk_tmp_min_datas = self.total_tmp_min_data[self.total_tmp_min_data[EMVar.time] == date_time]
                self.deliever(date_time, stk_tmp_min_datas)
        if self.mode == self.__class__.MODE_DAY:
            if not self.total_tmp_day_data.empty:
                self.deliever(date_time, self.total_tmp_day_data)
                pass



    def detect_stks(self, stks, date):
        self.registe_stks = list(set(self.registe_stks))
        if self.mode == self.__class__.MODE_MIN:
            tmp_total_min_data_pieces = []
            for stk in stks:
                if not stk in self.base_info_dic.keys():
                    model = sysBaseData.get_self_fund(stk)
                    self.base_info_dic[stk] = model.stkname
                    min_data = stk_utils.em_resample_tick_data(code=stk, date=date, need_insert_db=True)
                    if not min_data.empty:
                        tmp_total_min_data_pieces.append(min_data)

            if tmp_total_min_data_pieces:
                self.total_tmp_min_data = pd.concat(tmp_total_min_data_pieces, ignore_index=True)
                self.total_tmp_min_data.sort_values(EMVar.time, inplace=True)

        if self.mode == self.__class__.MODE_DAY:
            tmp_total_day_data_pieces = []
            for stk in stks:
                if not stk in self.base_info_dic.keys():
                    model = sysBaseData.get_self_fund(stk)
                    self.base_info_dic[stk] = model.stkname
                    model_data = model.get_index(start=self.start_date, indicator=['ma'], append_now=self.append_now)
                    # model_data = model.get_index(start='2004-01-01', indicator=['ma'], append_now=False, autype='hfq')
                    # model_data = model.get_index(start='2004-01-01', indicator=['ma'], append_now=False, autype=None)
                    if date in model_data.index:
                        dic_day_data = dict(model_data.ix[date])
                        day_data = pd.DataFrame(data=dic_day_data, index=[dic_day_data[EMVar.date]])
                        tmp_total_day_data_pieces.append(day_data)
            if tmp_total_day_data_pieces:
                self.total_tmp_day_data = pd.concat(tmp_total_day_data_pieces, ignore_index=True)





    def common_saile(self, assets, code, name , cur_close_price, cur_date_time,stop_win,stop_loss,is_reback, max_reback,max_keep_days,high_profit_reback,callback):
        baseModel = assets.get_own_stock(code)
        cur_date = time_utils.reformater(cur_date_time, formatter=EMVar.FORMATE_YMD)
        assets.update_stk_assets(stkcode=code, curprice=cur_close_price, cur_date=cur_date)
        if (cur_close_price - baseModel.trade_price) / baseModel.trade_price < stop_loss:
            print('止损卖出')
            assets.saile_stk(date_time=cur_date_time, stkcode=code, name=name,
                                  price=cur_close_price, reason='止损卖出', callback=callback)
            # assets.save_to_mongo()
            assets.update_one()
            return True
        elif (cur_close_price - baseModel.trade_price) / baseModel.trade_price > stop_win:
            print('止盈卖出')
            assets.saile_stk(date_time=cur_date_time, stkcode=code, name=name,
                                  price=cur_close_price, reason='止盈卖出', callback=callback)
            assets.update_one()
            return True
        if is_reback and baseModel.reback > max_reback and baseModel.high_profit > high_profit_reback:
            print('回撤卖出')
            assets.saile_stk(date_time=cur_date_time, stkcode=code, name=name,
                                  price=cur_close_price, reason='回撤卖出', callback=callback)
            assets.update_one()
            return True
        if baseModel.stk_keep_days >= max_keep_days:
            print('持仓日大于%d卖出' % max_keep_days)
            assets.saile_stk(date_time=cur_date_time, stkcode=code, name=name,
                                  price=cur_close_price, reason='持仓日大于%d卖出' % max_keep_days, callback=callback)
            assets.update_one()
            return True
        # else:
        #     print('UPDATE PRISE')
        #     cur_date = time_utils.reformater(cur_date_time, formatter=EMVar.FORMATE_YMD)
        #     assets.update_stk_assets(stkcode=code, curprice=cur_close_price,cur_date=cur_date)
        #
        #     return False





